import torch
import torch.optim as optim
from mnist.models import CNNMnist
from mnist.models import CNNDropoutMnist
from mnist.loader import load_mnist
import os
from torch import nn

def train_model(epochs=10):
    device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
    model = CNNMnist().to(device)
    optimizer = optim.Adam(model.parameters(), lr=0.001)
    train_loader, test_loader = load_mnist
    model.train()
    for epoch in range(epochs):
        for batch_idx, (data, target) in enumerate(train_loader):
            data, target = data.to(device), target.to(device)
            optimizer.zero_grad()
            output = model(data)
            loss = nn.functional.nll_loss(output, target)
            loss.backward()
            optimizer.step()

            if batch_idx % 100 == 0:
                print(f'Train Epoch: {epoch} [{batch_idx * len(data)}/{len(train_loader.dataset)}] Loss: {loss.item()}')
        # 在每个 epoch 结束后测试模型并打印准确率
        test_model(model, test_loader, device, epoch)

    print("训练测试完成，开始保存模型成功")
    os.makedirs('ml', exist_ok=True)
    torch.save(model.state_dict(), 'ml/mnist_cnn.pth')

def test_model(model, test_loader, device, epoch):
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            _, predicted = torch.max(output.data, 1)
            total += target.size(0)
            correct += (predicted == target).sum().item() # noqa

    accuracy = 100 * correct / total
    print(f'Epoch {epoch} Test Accuracy: {accuracy:.2f}%')


## 加入Dropout解决过拟合问题
def train_model_dropout():
    device = torch.device("mps" if torch.backends.mps.is_available() else "cpu")
    model = CNNDropoutMnist().to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    criterion = nn.CrossEntropyLoss()
    train_loader, test_loader = load_mnist()
    model.train()
    for epoch in range(10):
        for data, target in train_loader:
            data, target = data.to(device), target.to(device)
            optimizer.zero_grad()
            output = model(data)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
        # 在每个 epoch 结束后测试模型并打印准确率
        test_model_dropout(model, test_loader, device, epoch)

    print("训练测试完成，开始保存模型成功")
    torch.save(model.state_dict(), 'ml/mnist_cnn_dropout.pth')

def test_model_dropout(model, test_loader, device, epoch):
    model.eval()
    correct = 0
    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(device), target.to(device)
            output = model(data)
            pred = output.argmax(dim=1, keepdim=True)
            correct += pred.eq(target.view_as(pred)).sum().item()

    acc = 100. * correct / len(test_loader.dataset)
    print(f'Epoch {epoch}: Test Accuracy {acc:.2f}%')

if __name__ == '__main__':
    # train_model() # 训练模型，模型文件为mnist_cnn.pth
    train_model_dropout()  # 训练模型，模型文件为mnist_cnn_dropout.pth
